import numpy as np
from scipy.integrate import quad
from scipy.optimize import fsolve
import math
import pandas as pd

import matplotlib.pyplot as plt
import matplotlib.animation as animation

# 给定参数
p = 0.55  # 螺距，单位：米
b = p / (2 * np.pi)  # b的值
r_0_1 = 0.55 * 16  # 初始半径 r_0^(1)
theta_0_1 = 16 * 2 * np.pi  # 初始角度 theta_0^(1)


def binary_search(f, a, b, now, tol=1e-6, max_iter=1000):
    if f(a, now) * f(b, now) >= 0:
        raise ValueError("f(a) 和 f(b) 必须异号")

    for _ in range(max_iter):
        mid = (a + b) / 2
        f_mid = f(mid, now)
        if abs(f_mid) < tol:
            return mid
        elif f(a, now) * f_mid < 0:
            b = mid
        else:
            a = mid
    print("answer not find")
    return None


def polar_distance(theta1, theta2):
    r1 = b * theta1
    r2 = b * theta2
    # 计算角度差
    angle_diff = theta2 - theta1
    # 计算距离
    distance = math.sqrt((r1 - r2 * math.cos(angle_diff)) ** 2 + (r2 * math.sin(angle_diff)) ** 2)
    return distance


# 示例：求解方程 x**2 - 4 = 0 在区间 [1, 3] 内的根
def cal_body(theta, now):
    length = 1.65
    distance = polar_distance(theta, now)
    return distance - length


def cal_head(theta, now):
    length = 2.86
    distance = polar_distance(theta, now)
    return distance - length


def get_theta_position(now):  # 给龙头，给全部
    result = [now]
    num = 224 - 1 - 1
    body = binary_search(cal_head, now + 0.001, now + 1, now)
    result.append(body)
    for i in range(num):
        next = binary_search(cal_body, body + 0.001, body + 1, body)
        if next > 16 * 2 * np.pi:
            result.append(0)
        else:
            result.append(next)
        body = next
    return result


def integrand(theta):
    return np.sqrt((b * theta) ** 2 + b ** 2)


# 定义求解 theta 的方程
def find_theta_1(theta_1, theta_0, length=1):
    return quad(integrand, theta_1, theta_0)[0] - length


def get_theta_time(now, time_now, time_next, timespace):  # 给一点，算时序
    result = [now]
    t = 0
    while t < time_next - time_now:
        t += timespace
        theta_prev = result[-1]
        theta_new = fsolve(find_theta_1, theta_prev - 1, args=theta_prev)[0]
        result.append(theta_new)
    return result


def cal_total():
    head = get_theta_time(theta_0_1, 0, 300, 1)
    ans = []
    for i in head:
        tem = get_theta_position(i)
        ans.append(tem)
    return ans


def calculate_slope(theta):
    up = (np.sin(theta) + theta * np.cos(theta))
    down = (np.cos(theta) - theta * np.sin(theta))
    return np.arctan(up / down)


def cal_velocity_components(theta_a, theta_b, v_a):
    """计算速度分解到ab线段方向的分量"""
    # 计算螺线在a点和b点的切线角度
    r1 = b * theta_a
    r2 = b * theta_b
    x_0 = r1 * np.cos(theta_a)
    y_0 = r1 * np.sin(theta_a)
    x_b = r2 * np.cos(theta_b)
    y_b = r2 * np.sin(theta_b)
    tangent_angle_a = calculate_slope(theta_a)
    tangent_angle_b = calculate_slope(theta_b)

    # 计算ab线段的角度
    ab_angle = np.arctan2(y_b - y_0, x_b - x_0)

    # 计算速度分解到ab线段方向的分量
    v_a_line = v_a * np.cos(tangent_angle_a - ab_angle)
    v_b_line = v_a_line
    v_b = v_b_line / np.cos(tangent_angle_b - ab_angle)
    return abs(v_b)


def transpose_2d(data):
    transposed = list(map(list, zip(*data)))
    return transposed


def output():
    with pd.ExcelWriter('result1.xlsx') as writer:
        # 算极角
        head = get_theta_time(theta_0_1, 0, 300, 1)  # 300s的龙头极角变化
        theta_all = cal_total()  # 每时刻，每个点的极角
        print("完成极角计算")

        # 算坐标
        pos = []
        for i in theta_all:
            tem = []
            for j in i:
                r = b * j
                x = r * np.cos(j)
                y = r * np.sin(j)
                tem.extend([x, y])
            pos.append(tem)
        df = pd.DataFrame(pos).transpose()
        index = ["龙头x (m)", "龙头y (m)"]
        columns = [f"{i} s" for i in range(301)]
        for i in range(1, 222):
            index.extend([f"第{i}节龙身x (m)", f"第{i}节龙身y (m)"])
        index.extend(["龙尾x (m)", "龙尾y (m)", "龙尾（后）y (m)", "龙尾（后）y (m)"])
        df.columns = columns
        df.index = index
        df.to_excel(writer, sheet_name="坐标", engine="openpyxl")
        print("完成坐标计算")
        # 算速度
        vel = [[1] for i in range(0, 301)]
        for time in range(0, 301):  # [0,300]时间
            for point in range(1, 224):  # [0,223]个点
                theta_a = theta_all[time][point - 1]
                theta_b = theta_all[time][point]
                v_a = vel[time][point - 1]
                if theta_b == 0:
                    vel[time].append(0)
                else:
                    vel[time].append(cal_velocity_components(theta_a, theta_b, v_a))
        vel = transpose_2d(vel)
        index = ["龙头 (m/s)"]
        for i in range(1, 222):
            index.append(f"第{i}节龙身 (m/s)")
        index.extend(["龙尾 (m/s)", "龙尾（后） (m/s)"])
        df = pd.DataFrame(vel, index=index, columns=columns)
        df.to_excel(writer, sheet_name="速度", engine="openpyxl")
        print("完成速度计算")


def draw():
    # 计算等距螺线的坐标
    theta_values = np.linspace(0, theta_0_1, 1000)
    r_values = b * theta_values
    x_values = r_values * np.cos(theta_values)
    y_values = r_values * np.sin(theta_values)

    # 创建动画
    fig, ax = plt.subplots()
    ax.set_xlim(-r_0_1, r_0_1)
    ax.set_ylim(-r_0_1, r_0_1)
    ax.set_aspect('equal')
    ax.plot(x_values, y_values, 'b-', label='line')  # 绘制等距螺线
    lines = [ax.plot([], [], 'ro-')[0] for _ in range(3)]  # 创建三个线段

    def init():
        for line in lines:
            line.set_data([], [])
        return lines

    def animate(i):
        theta1 = theta_values[i]
        theta2 = theta1 + 2.86 / (b * theta1)  # 计算第一个线段的另一端的角度
        theta3 = theta2 + 1.65 / (b * theta2)  # 计算第二个线段的另一端的角度
        theta4 = theta3 + 1.65 / (b * theta3)  # 计算第三个线段的另一端的角度

        r1 = b * theta1
        r2 = b * theta2
        r3 = b * theta3
        r4 = b * theta4

        x1 = r1 * np.cos(theta1)
        y1 = r1 * np.sin(theta1)
        x2 = r2 * np.cos(theta2)
        y2 = r2 * np.sin(theta2)
        x3 = r3 * np.cos(theta3)
        y3 = r3 * np.sin(theta3)
        x4 = r4 * np.cos(theta4)
        y4 = r4 * np.sin(theta4)

        lines[0].set_data([x1, x2], [y1, y2])
        lines[1].set_data([x2, x3], [y2, y3])
        lines[2].set_data([x3, x4], [y3, y4])

        return lines

    ani = animation.FuncAnimation(fig, animate, frames=len(theta_values), init_func=init, blit=True, interval=50)

    plt.legend()
    plt.show()
    ani.save("ttt.gif", "pillow")


if __name__ == "__main__":
    # output() # 输出result1.xlsx
    draw()
